# -*- coding: utf-8 -*-

import tushare as ts
import numpy as np
import os
import pandas as pd

tushare_dir = '../dataset/tushare'
if os.path.exists(tushare_dir) == False:
    os.mkdir(tushare_dir)

stock = pd.DataFrame()
if os.path.exists(tushare_dir + '/stock.csv') == False:
    stock = ts.get_stock_basics()
    stock.to_csv(tushare_dir + '/stock.csv')
else:
    stock=stock.from_csv(tushare_dir + '/stock.csv')

def save_hist():
    hist_dir = tushare_dir+ '/hist'
    if os.path.exists(hist_dir) == False:
        os.mkdir(hist_dir)
    for code in stock.index:
        code = str(code)
        print (code)
        hist = ts.get_hist_data(code)
        if hist is not None:
            hist.to_csv(hist_dir + '/' + code + '.csv')


def get_all_h_data(code):
    end_data = '2017-12-16'
    date = stock.loc[int(code),'timeToMarket']
    date = str(date)
    if len(date) != 8:
        return None
    date_ = date[0:4] + '-' + date[4:6] + '-' + date[6:8]
    return ts.get_h_data(code, start=date_, end=end_data)

def save_h():
    h_dir = tushare_dir+'/h'
    if os.path.exists(h_dir) == False:
        os.mkdir(h_dir)
    for code in stock.index:
        try:
            code = ('000000' + str(code))[-6:]
            p = h_dir + '/' + code + '.csv'
            if os.path.exists(p):
                print (code)
            else:
                print (code + '+')
                h = get_all_h_data(code)
                if h is not None:
                    print ('ok')
                    h.to_csv(p)
                else:
                    print ('bad')
        except Exception as e:
            print (e.message)
            continue
save_h()
pass